package compiler;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;

import javax.tools.DiagnosticCollector;
import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.SimpleJavaFileObject;
import javax.tools.StandardLocation;
import javax.tools.ToolProvider;
import javax.tools.JavaCompiler.CompilationTask;
import javax.tools.JavaFileObject.Kind;

public class ScriptCompiler<T> {
	static final String JAVA_EXTENSION = ".java";

	private final ScriptClassLoader classLoader;
	private final JavaCompiler compiler;
	private DiagnosticCollector<JavaFileObject> diagnostics;
	private final ScriptFileManager javaFileManager;

	public ScriptCompiler(ClassLoader loader) {
		compiler = ToolProvider.getSystemJavaCompiler();
		if (compiler == null) {
			throw new IllegalStateException("Cannot find the system Java compiler.");
		}
		classLoader = new ScriptClassLoader(loader);
		diagnostics = new DiagnosticCollector<JavaFileObject>();
		final JavaFileManager fileManager = compiler.getStandardFileManager(diagnostics, null, null);
		javaFileManager = new ScriptFileManager(fileManager, classLoader);
	}

	public synchronized Class<T> compile(final String qualifiedClassName, final CharSequence javaSource,
			final DiagnosticCollector<JavaFileObject> diagnosticsList, final Class<?>... types) 
			throws ScriptCompilerException {
		if (diagnosticsList != null) {
			diagnostics = diagnosticsList;
		} else {
			diagnostics = new DiagnosticCollector<JavaFileObject>();
		}
		Map<String, CharSequence> classes = new HashMap<String, CharSequence>(1);
		classes.put(qualifiedClassName, javaSource);
		Map<String, Class<T>> compiled = compile(classes, diagnosticsList);
		Class<T> newClass = compiled.get(qualifiedClassName);
		return castable(newClass, types);
	}

	public synchronized Map<String, Class<T>> compile(final Map<String, CharSequence> classes,
			final DiagnosticCollector<JavaFileObject> diagnosticsList) throws ScriptCompilerException {
		List<JavaFileObject> sources = new ArrayList<JavaFileObject>();
		for (Entry<String, CharSequence> entry : classes.entrySet()) {
			String qualifiedClassName = entry.getKey();
			CharSequence javaSource = entry.getValue();
			if (javaSource != null) {
				final int dotPos = qualifiedClassName.lastIndexOf('.');
				final String className = dotPos == -1 ? qualifiedClassName : qualifiedClassName.substring(dotPos + 1);
				final String packageName = dotPos == -1 ? "" : qualifiedClassName.substring(0, dotPos);
				final ScriptJavaFileObject source = new ScriptJavaFileObject(className, javaSource);
				sources.add(source);
				
				javaFileManager.putFileForInput(StandardLocation.SOURCE_PATH, packageName,
						className + JAVA_EXTENSION, source);
			}
		}
		
		final CompilationTask task = compiler.getTask(null, javaFileManager, diagnostics, null, null, sources);
		final Boolean result = task.call();
		if (result == null || !result.booleanValue()) {
			throw new ScriptCompilerException("Compilation failed.", classes.keySet(), diagnostics);
		}
		try {
			Map<String, Class<T>> compiled = new HashMap<String, Class<T>>();
			for (String qualifiedClassName : classes.keySet()) {
				final Class<T> newClass = loadClass(qualifiedClassName);
				compiled.put(qualifiedClassName, newClass);
			}
			return compiled;
		} catch (ClassNotFoundException e) {
			throw new ScriptCompilerException(classes.keySet(), e, diagnostics);
		} catch (IllegalArgumentException e) {
			throw new ScriptCompilerException(classes.keySet(), e, diagnostics);
		} catch (SecurityException e) {
			throw new ScriptCompilerException(classes.keySet(), e, diagnostics);
		}
	}

	public Class<T> loadClass(final String qualifiedClassName) throws ClassNotFoundException {
		return (Class<T>)classLoader.loadClass(qualifiedClassName);
	}

	private Class<T> castable(Class<T> newClass, Class<?>... types) throws ClassCastException {
		for (Class<?> type : types) {
			if (!type.isAssignableFrom(newClass)) {
				throw new ClassCastException(type.getName());
			}
		}
		return newClass;
	}

	static URI toURI(String name) {
		try {
			return new URI(name);
		} catch (URISyntaxException e) {
			throw new RuntimeException(e);
		}
	}

	public ClassLoader getClassLoader() {
		return javaFileManager.getClassLoader();
	}
}

/**
 * A JavaFileManager which manages Java source and classes. This FileManager
 * delegates to the JavaFileManager and the ClassLoaderImpl provided in the
 * constructor. The sources are all in memory CharSequence instances and the
 * classes are all in memory byte arrays.
 */
final class ScriptFileManager extends ForwardingJavaFileManager<JavaFileManager> {
	// the delegating class loader (passed to the constructor)
	private final ScriptClassLoader classLoader;

	// Internal map of filename URIs to JavaFileObjects.
	private final Map<URI, JavaFileObject> fileObjects = new HashMap<URI, JavaFileObject>();

	public ScriptFileManager(JavaFileManager fileManager, ScriptClassLoader classLoader) {
		super(fileManager);
		this.classLoader = classLoader;
	}

	public ClassLoader getClassLoader() {
		return classLoader;
	}

	@Override
	public FileObject getFileForInput(Location location, String packageName, String relativeName) throws IOException {
		FileObject o = fileObjects.get(uri(location, packageName, relativeName));
		if (o != null) {
			return o;
		} else {
			return super.getFileForInput(location, packageName, relativeName);
		}
	}

	public void putFileForInput(StandardLocation location, String packageName, String relativeName, JavaFileObject file) {
		fileObjects.put(uri(location, packageName, relativeName), file);
	}

	private URI uri(Location location, String packageName, String relativeName) {
		return ScriptCompiler.toURI(location.getName() + '/' + packageName + '/' + relativeName);
	}

	@Override
	public JavaFileObject getJavaFileForOutput(Location location, String qualifiedName, Kind kind, 
			FileObject outputFile) throws IOException {
		JavaFileObject file = new ScriptJavaFileObject(qualifiedName, kind);
		classLoader.add(qualifiedName, file);
		return file;
	}

	@Override
	public ClassLoader getClassLoader(JavaFileManager.Location location) {
		return classLoader;
	}

	@Override
	public String inferBinaryName(Location loc, JavaFileObject file) {
		String result;
		if (file instanceof ScriptJavaFileObject) {
			result = file.getName();
		} else {
			result = super.inferBinaryName(loc, file);
		}
		return result;
	}

	@Override
	public Iterable<JavaFileObject> list(Location location, String packageName, Set<Kind> kinds, boolean recurse) 
	throws IOException {
		Iterable<JavaFileObject> result = super.list(location, packageName, kinds,
				recurse);
		ArrayList<JavaFileObject> files = new ArrayList<JavaFileObject>();
		if (location == StandardLocation.CLASS_PATH
				&& kinds.contains(JavaFileObject.Kind.CLASS)) {
			for (JavaFileObject file : fileObjects.values()) {
				if (file.getKind() == Kind.CLASS && file.getName().startsWith(packageName))
					files.add(file);
			}
			files.addAll(classLoader.files());
		} else if (location == StandardLocation.SOURCE_PATH
				&& kinds.contains(JavaFileObject.Kind.SOURCE)) {
			for (JavaFileObject file : fileObjects.values()) {
				if (file.getKind() == Kind.SOURCE && file.getName().startsWith(packageName))
					files.add(file);
			}
		}
		for (JavaFileObject file : result) {
			files.add(file);
		}
		return files;
	}
}

final class ScriptJavaFileObject extends SimpleJavaFileObject {
	// If kind == CLASS, this stores byte code from openOutputStream
	private ByteArrayOutputStream byteCode;

	// if kind == SOURCE, this contains the source text
	private final CharSequence source;
	
	ScriptJavaFileObject(final String baseName, final CharSequence source) {
		super(ScriptCompiler.toURI(baseName + ScriptCompiler.JAVA_EXTENSION),
				Kind.SOURCE);
		this.source = source;
	}
	
	ScriptJavaFileObject(final String name, final Kind kind) {
		super(ScriptCompiler.toURI(name), kind);
		this.source = null;
	}

	/**
	 * Return the source code content
	 * 
	 * @see javax.tools.SimpleJavaFileObject#getCharContent(boolean)
	 */
	@Override
	public CharSequence getCharContent(final boolean ignoreEncodingErrors) throws UnsupportedOperationException {
		if (source == null) {
			throw new NullPointerException();
		}
		return source;
	}
	
	@Override
	public InputStream openInputStream() {
		return new ByteArrayInputStream(getByteCode());
	}
	
	@Override
	public OutputStream openOutputStream() {
		byteCode = new ByteArrayOutputStream();
		return byteCode;
	}

	
	byte[] getByteCode() {
		return byteCode.toByteArray();
	}
}

final class ScriptClassLoader extends ClassLoader {
	private final Map<String, JavaFileObject> classes = new HashMap<String, JavaFileObject>();

	ScriptClassLoader(final ClassLoader parentClassLoader) {
		super(parentClassLoader);
	}
	
	Collection<JavaFileObject> files() {
		return Collections.unmodifiableCollection(classes.values());
	}

	@Override
	protected Class<?> findClass(final String qualifiedClassName) throws ClassNotFoundException {
		JavaFileObject file = classes.get(qualifiedClassName);
		if (file != null) {
			byte[] bytes = ((ScriptJavaFileObject)file).getByteCode();
			return defineClass(qualifiedClassName, bytes, 0, bytes.length);
		} else {
			// Workaround for "feature" in Java 6
			// see http://bugs.sun.com/bugdatabase/view_bug.do?bug_id=6434149
			try {
				Class<?> c = Class.forName(qualifiedClassName);
				return c;
			} catch (ClassNotFoundException nf) {
				// Ignore and fall through
			}
			return super.findClass(qualifiedClassName);
		}
	}
	
	void add(final String qualifiedClassName, final JavaFileObject javaFile) {
		classes.put(qualifiedClassName, javaFile);
	}

	@Override
	protected synchronized Class<?> loadClass(final String name, final boolean resolve) throws ClassNotFoundException {
		return super.loadClass(name, resolve);
	}

	@Override
	public InputStream getResourceAsStream(final String name) {
		if (name.endsWith(".class")) {
			String qualifiedClassName = name.substring(0, name.length() - ".class".length()).replace('/', '.');
			ScriptJavaFileObject file = (ScriptJavaFileObject)classes.get(qualifiedClassName);
			if (file != null) {
				return new ByteArrayInputStream(file.getByteCode());
			}
		}
		return super.getResourceAsStream(name);
	}
}
